{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PyTorch Lightning for MNIST\n",
    "\n",
    "Lighting to easily train and test a model on MNIST dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: lightning in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (2.4.0)\n",
      "Collecting lightning\n",
      "  Downloading lightning-2.5.5-py3-none-any.whl.metadata (39 kB)\n",
      "Requirement already satisfied: torchmetrics in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (1.7.3)\n",
      "Collecting torchmetrics\n",
      "  Downloading torchmetrics-1.8.2-py3-none-any.whl.metadata (22 kB)\n",
      "Requirement already satisfied: PyYAML<8.0,>5.4 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from lightning) (6.0.2)\n",
      "Requirement already satisfied: fsspec<2027.0,>=2022.5.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from fsspec[http]<2027.0,>=2022.5.0->lightning) (2024.12.0)\n",
      "Requirement already satisfied: lightning-utilities<2.0,>=0.10.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from lightning) (0.14.3)\n",
      "Requirement already satisfied: packaging<27.0,>=20.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from lightning) (24.1)\n",
      "Requirement already satisfied: torch<4.0,>=2.1.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from lightning) (2.2.2)\n",
      "Requirement already satisfied: tqdm<6.0,>=4.57.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from lightning) (4.67.1)\n",
      "Requirement already satisfied: typing-extensions<6.0,>4.5.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from lightning) (4.14.1)\n",
      "Requirement already satisfied: pytorch-lightning in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from lightning) (2.5.2)\n",
      "Requirement already satisfied: numpy>1.20.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from torchmetrics) (1.26.4)\n",
      "Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from fsspec[http]<2027.0,>=2022.5.0->lightning) (3.12.14)\n",
      "Requirement already satisfied: setuptools in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from lightning-utilities<2.0,>=0.10.0->lightning) (75.1.0)\n",
      "Requirement already satisfied: filelock in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from torch<4.0,>=2.1.0->lightning) (3.16.1)\n",
      "Requirement already satisfied: sympy in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from torch<4.0,>=2.1.0->lightning) (1.13.3)\n",
      "Requirement already satisfied: networkx in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from torch<4.0,>=2.1.0->lightning) (3.4.2)\n",
      "Requirement already satisfied: jinja2 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from torch<4.0,>=2.1.0->lightning) (3.1.4)\n",
      "Requirement already satisfied: aiohappyeyeballs>=2.5.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2027.0,>=2022.5.0->lightning) (2.6.1)\n",
      "Requirement already satisfied: aiosignal>=1.4.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2027.0,>=2022.5.0->lightning) (1.4.0)\n",
      "Requirement already satisfied: async-timeout<6.0,>=4.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2027.0,>=2022.5.0->lightning) (5.0.1)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2027.0,>=2022.5.0->lightning) (24.2.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2027.0,>=2022.5.0->lightning) (1.5.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2027.0,>=2022.5.0->lightning) (6.1.0)\n",
      "Requirement already satisfied: propcache>=0.2.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2027.0,>=2022.5.0->lightning) (0.2.1)\n",
      "Requirement already satisfied: yarl<2.0,>=1.17.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2027.0,>=2022.5.0->lightning) (1.18.3)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from jinja2->torch<4.0,>=2.1.0->lightning) (3.0.2)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from sympy->torch<4.0,>=2.1.0->lightning) (1.3.0)\n",
      "Requirement already satisfied: idna>=2.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from yarl<2.0,>=1.17.0->aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2027.0,>=2022.5.0->lightning) (3.10)\n",
      "Downloading lightning-2.5.5-py3-none-any.whl (828 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m828.5/828.5 kB\u001b[0m \u001b[31m3.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hDownloading torchmetrics-1.8.2-py3-none-any.whl (983 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m983.2/983.2 kB\u001b[0m \u001b[31m20.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: torchmetrics, lightning\n",
      "  Attempting uninstall: torchmetrics\n",
      "    Found existing installation: torchmetrics 1.7.3\n",
      "    Uninstalling torchmetrics-1.7.3:\n",
      "      Successfully uninstalled torchmetrics-1.7.3\n",
      "  Attempting uninstall: lightning\n",
      "    Found existing installation: lightning 2.4.0\n",
      "    Uninstalling lightning-2.4.0:\n",
      "      Successfully uninstalled lightning-2.4.0\n",
      "Successfully installed lightning-2.5.5 torchmetrics-1.8.2\n",
      "Requirement already satisfied: torch in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (2.2.2)\n",
      "Requirement already satisfied: torchvision in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (0.17.2)\n",
      "Requirement already satisfied: filelock in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from torch) (3.16.1)\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from torch) (4.14.1)\n",
      "Requirement already satisfied: sympy in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from torch) (1.13.3)\n",
      "Requirement already satisfied: networkx in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from torch) (3.4.2)\n",
      "Requirement already satisfied: jinja2 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from torch) (3.1.4)\n",
      "Requirement already satisfied: fsspec in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from torch) (2024.12.0)\n",
      "Requirement already satisfied: numpy in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from torchvision) (1.26.4)\n",
      "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from torchvision) (11.0.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from jinja2->torch) (3.0.2)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /Users/rowel/anaconda3/envs/agents/lib/python3.10/site-packages (from sympy->torch) (1.3.0)\n"
     ]
    }
   ],
   "source": [
    "! pip install lightning torchmetrics --upgrade\n",
    "! pip install torch torchvision --upgrade"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### DataModule for MNIST\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import lightning as L\n",
    "import torch.nn as nn\n",
    "from torchvision import transforms\n",
    "\n",
    "# create an MNIST datamodule\n",
    "\n",
    "class MNISTDataModule(L.LightningDataModule):\n",
    "    def __init__(self, batch_size: int = 32):\n",
    "        super().__init__()\n",
    "        self.batch_size = batch_size\n",
    "\n",
    "    def prepare_data(self):\n",
    "        torchvision.datasets.MNIST(root=\"./data\", download=True)\n",
    "\n",
    "    def setup(self, stage=None):\n",
    "        if stage == \"fit\" or stage is None:\n",
    "            self.mnist_train = torchvision.datasets.MNIST(\n",
    "                root=\"./data\", train=True, download=True, transform=transforms.ToTensor()\n",
    "            )\n",
    "        if stage == \"test\" or stage is None:\n",
    "            self.mnist_test = torchvision.datasets.MNIST(\n",
    "                root=\"./data\", train=False,download=True, transform=transforms.ToTensor()\n",
    "            )\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return torch.utils.data.DataLoader(\n",
    "            self.mnist_train, batch_size=self.batch_size, shuffle=True\n",
    "        )\n",
    "\n",
    "    def test_dataloader(self):\n",
    "        return torch.utils.data.DataLoader(\n",
    "            self.mnist_test, batch_size=self.batch_size, shuffle=False\n",
    "        )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model for MNIST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3-layer CNN model\n",
    "class CNN(nn.Module):\n",
    "    def __init__(self, num_classes=10):\n",
    "        super(CNN, self).__init__()\n",
    "        self.features = nn.Sequential(\n",
    "            nn.Conv2d(1, 16, kernel_size=3),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size=2),\n",
    "            nn.Conv2d(16, 32, kernel_size=3),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size=2),\n",
    "        )\n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.Linear(800, 128),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(128, num_classes)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        y = self.features(x)\n",
    "        y = torch.flatten(y, 1)\n",
    "        y = self.classifier(y)\n",
    "        \n",
    "        return y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LightningModule for MNIST\n",
    "\n",
    "Notice this module uses the sample `CNN()` model defined above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build a pytorch lightning model\n",
    "from torchmetrics import Accuracy\n",
    "\n",
    "class LitModel(L.LightningModule):\n",
    "    def __init__(self, model=CNN()):\n",
    "        super().__init__()\n",
    "        self.model = model\n",
    "        self.loss_fn = nn.CrossEntropyLoss()\n",
    "        self.accuracy = Accuracy(task=\"multiclass\", num_classes=10)\n",
    "        self.avg_acc = []\n",
    "        self.avg_loss = []\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.model(x)\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        y_hat = self(x)\n",
    "        loss = self.loss_fn(y_hat, y)\n",
    "        self.log(\"train_loss\", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True)\n",
    "        return loss\n",
    "\n",
    "    def on_train_epoch_end(self):\n",
    "        lr = self.trainer.optimizers[0].param_groups[0][\"lr\"]\n",
    "        self.log(\"learning_rate\", lr, on_step=False, on_epoch=True, prog_bar=True, logger=True)    \n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)\n",
    "        # linear decay\n",
    "        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer,\n",
    "                                                      lr_lambda=lambda epoch: 1 - epoch / self.trainer.max_epochs)\n",
    "        return [optimizer], [scheduler]\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        y_hat = self(x)\n",
    "        loss = self.loss_fn(y_hat, y)\n",
    "        acc = self.accuracy(y_hat, y)\n",
    "        self.avg_acc.append(acc)\n",
    "        self.avg_loss.append(loss)\n",
    "        return loss, acc\n",
    "\n",
    "    def on_test_epoch_end(self):\n",
    "        loss = torch.stack(self.avg_loss).mean()\n",
    "        acc = torch.stack(self.avg_acc).mean()\n",
    "        self.log(\"test_loss\", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True)\n",
    "        self.log(\"test_acc\", acc, on_step=False, on_epoch=True, prog_bar=True, logger=True)\n",
    "        self.avg_acc = []\n",
    "        self.avg_loss = []\n",
    "        return loss, acc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training and Testing the Model\n",
    "\n",
    "Note how we use the `Trainer` class to train and test the model. Also, note the selection of accelerator and precision."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.\n",
      "GPU available: False, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "HPU available: False, using: 0 HPUs\n",
      "\n",
      "  | Name     | Type               | Params | Mode \n",
      "--------------------------------------------------------\n",
      "0 | model    | CNN                | 108 K  | train\n",
      "1 | loss_fn  | CrossEntropyLoss   | 0      | train\n",
      "2 | accuracy | MulticlassAccuracy | 0      | train\n",
      "--------------------------------------------------------\n",
      "108 K     Trainable params\n",
      "0         Non-trainable params\n",
      "108 K     Total params\n",
      "0.434     Total estimated model params size (MB)\n",
      "14        Modules in train mode\n",
      "0         Modules in eval mode\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 9: 100%|██████████| 1875/1875 [00:33<00:00, 56.12it/s, v_num=1, train_loss=0.00237, learning_rate=0.0001]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=10` reached.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 9: 100%|██████████| 1875/1875 [00:33<00:00, 56.09it/s, v_num=1, train_loss=0.00237, learning_rate=0.0001]\n",
      "Testing DataLoader 0: 100%|██████████| 313/313 [00:03<00:00, 100.85it/s]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\">        Test metric        </span>┃<span style=\"font-weight: bold\">       DataLoader 0        </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">         test_acc          </span>│<span style=\"color: #800080; text-decoration-color: #800080\">    0.9915135502815247     </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">         test_loss         </span>│<span style=\"color: #800080; text-decoration-color: #800080\">   0.030096787959337234    </span>│\n",
       "└───────────────────────────┴───────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1m       Test metric       \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m      DataLoader 0       \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│\u001b[36m \u001b[0m\u001b[36m        test_acc         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m   0.9915135502815247    \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36m        test_loss        \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m  0.030096787959337234   \u001b[0m\u001b[35m \u001b[0m│\n",
       "└───────────────────────────┴───────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[{'test_loss': 0.030096787959337234, 'test_acc': 0.9915135502815247}]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# train for 10 epochs\n",
    "\n",
    "dm = MNISTDataModule()\n",
    "model = LitModel()\n",
    "trainer = L.Trainer(max_epochs=10, accelerator=\"auto\")\n",
    "trainer.fit(model, dm)\n",
    "trainer.test(model, dm)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "agents",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
